import { ChromaClient } from 'chromadb'
import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
import { Chroma } from "@langchain/community/vectorstores/chroma";

const listCollections = async () =>{
    const client = new ChromaClient({
        path: "http://117.72.38.226:8000"
      });
    const colls =await client.listCollections()
    for (let i = 0; i < colls.length; i++) {
        console.log(colls[i].name)
    }
}

// 清空数据集
const clearVectorDB =async (collName)=>{
    const client = new ChromaClient({
        path: "http://117.72.38.226:8000"
      });
    if( collName !="" ){
        await client.deleteCollection({name:collName})
        console.log("删除数据集：",collName)
    }else{
        const colls =await client.listCollections()
        console.log("现有数据集：",colls.length)
        for (let i = 0; i < colls.length; i++) {
            await client.deleteCollection({
                name: colls[i].name
            })
        }
        console.log("删除全部数据集")
    }
}

 // 创建数据集
const createCollection = async (collName)=>{
    if (collName =="") {
        console.error("还未指定数据集")
        return
    }
    const client = new ChromaClient({
        path: "http://117.72.38.226:8000"
      });
    const coll  =await client.getOrCreateCollection({name:collName})
    console.log("现有数据条数：",await coll.count())
}
// 挑选数据
const peekData = async (collName)=>{
    if (collName =="") {
        console.error("还未指定数据集")
        return
    }
    const client = new ChromaClient({
        path: "http://117.72.38.226:8000"
      });
    const coll  =await client.getCollection({name:collName})
    const result=await coll.peek({
            limit: 10,
        })
    console.log(result)
}

const selectData = async (collName)=>{
    if (collName =="") {
        console.error("还未指定数据集")
        return
    }
    const embeddings = new OllamaEmbeddings({
        model: "milkey/m3e", // default value
        baseUrl: "http://117.72.38.226:11434", // default value
        requestOptions: {
        useMMap: true,
        numThread: 6,
        // numGpu: 1,
        },
    });
    const vectorStore = new Chroma(
        embeddings,
        {
        collectionName: collName,
        url: "http://117.72.38.226:8000",
        }
    )
    const result = await vectorStore.similaritySearch("愚公做了什么", 3);
    console.log(result)
}

listCollections()
// clearVectorDB("go-chroma")
// createCollection("hcy-test1")
// peekData("go-chroma")
// selectData("go-chroma")
